# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/16 13:17 
@Author  : ZhangShenao 
@File    : 3.使用RouterChain路由链.py 
@Desc    : 使用RouterChain路由链,组成一个完整的决策系统
"""
import os

import dotenv
from langchain.chains.conversation.base import ConversationChain
from langchain.chains.router import MultiPromptChain
from langchain.chains.router.llm_router import RouterOutputParser, LLMRouterChain
from langchain.chains.router.multi_prompt_prompt import MULTI_PROMPT_ROUTER_TEMPLATE as RouterTemplate
from langchain_core.prompts import PromptTemplate
from langchain_openai import OpenAI

# 加载环境变量
dotenv.load_dotenv()

# 针对两个不同的场景,创建对应的PromptTemplate
flower_care_template = '''
你是一个经验丰富的园丁，擅长解答关于养花育花的问题。
下面是需要你来回答的问题:
{input}
'''
flower_decorate_template = '''
你是一位网红插花大师，擅长解答关于鲜花装饰的问题。
下面是需要你来回答的问题:
{input}
'''

# 构造Prompt列表
prompt_infos = [
    {
        'key': 'flower_care',
        'description': '适合回答关于鲜花护理的问题',
        'template': flower_care_template,
    },
    {
        'key': 'flower_decorate',
        'description': '适合回答关于鲜花装饰的问题',
        'template': flower_decorate_template,
    }
]

# 构造LLM
llm = OpenAI(model_name='gpt-3.5-turbo', openai_api_base=os.getenv('OPENAI_API_BASE'))

# 遍历Prompt列表,构造对应的Chain
chains = {}
for prompt_info in prompt_infos:
    prompt = PromptTemplate(template=prompt_info['template'], input_variables=['input'])
    chain = prompt | llm
    chains[prompt_info['key']] = chain

# 构建路由链
destinations = [f'{p['key']}：{p['description']}' for p in prompt_infos]
router_template = RouterTemplate.format(destinations='\n'.join(destinations))
router_prompt = PromptTemplate(
    template=router_template,
    input_variables=['input'],
    output_parser=RouterOutputParser(),
)
print(f'router_prompt={router_prompt}')
router_chain = LLMRouterChain.from_llm(llm=llm, prompt=router_prompt, verbose=True)

# 构建默认链
default_chain = ConversationChain(
    llm=llm,
    output_key="text",
    verbose=True
)

# 使用MultiPromptChain,将多个链整合到一起,实现路由功能
# MultiPromptChain类是一个多路选择链,它使用一个LLM路由器链在多个提示之间进行选择
chain = MultiPromptChain(
    router_chain=router_chain,
    destination_chains=chains,
    default_chain=default_chain,
    verbose=True
)

# 执行Chain
print(chain.invoke({'input': '如何为玫瑰花浇水？'}))
